# 基于XGBoost原生接口的分类
from sklearn.datasets import load_iris
import xgboost as xgb
from xgboost import plot_importance
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, multilabel_confusion_matrix
import pandas as pd
import numpy as np
import graphviz

data = pd.read_csv('D:/lung_cancer/data/data.csv')
X = []
y = []
for i in range(len(data)):
    one_feature = [data['z'][i], data['x'][i], data['y'][i], data['r'][i],
                   data['patientWeight'][i], data['patientSex'][i], data['patientAge'][i],
                   data['part_suvmax'][i], data['suv_avg'][i], data['suv_std'][i]]
    X.append(one_feature)
    y.append(data['cancer_type'][i]-1)

X_train, X_test, y_train, y_test = train_test_split(np.array(X, dtype=np.float), np.array(y, dtype=np.int), test_size=0.3, random_state=1234565)

params = {
    'booster': 'gbtree',
    'objective': 'multi:softmax',
    'num_class': 5,
    'gamma': 0.1,
    'max_depth': 6,
    'lambda': 2,
    'subsample': 0.7,
    'colsample_bytree': 0.7,
    'min_child_weight': 3,
    'silent': 1,
    'eta': 0.1,
    'seed': 1000,
    'nthread': 4
}

plst = params.items()

dtrain = xgb.DMatrix(X_train, y_train)
num_rounds = 500
model = xgb.train(params, dtrain, num_rounds)

# 可视化模型
img1 = xgb.to_graphviz(model, num_trees=0)
img1.format = 'png'
img1.view('./img1')


# 保存模型
model.save_model('test.model')
# 导出模型
model.dump_model('model.txt')

# 对测试集进行测试
dtest = xgb.DMatrix(X_test)
ans = model.predict(dtest)

print('ans: ', ans)
print(type(ans))

ans2 = model.predict(dtest, pred_contribs=True)
print('ans2: ', ans2)

print('y_test: ', y_test)
print(type(y_test))

target_names = ['1', '2', '3', '4', '0']
result_statis = classification_report(y_test, ans, target_names=target_names)

print(result_statis)

# 计算准确率
cnt1 = 0
cnt2 = 0
for i in range(len(y_test)):
    if ans[i] == y_test[i]:
        cnt1 += 1
    else:
        cnt2 += 1

train_pred = model.predict(dtrain)
train_accuracy = accuracy_score(y_train, train_pred)
print("test accuracy: %.2f %%" % (100*cnt1/(cnt1+cnt2)))

print("train accuracy: %.2f %%" % (train_accuracy*100))


# 显示特征重要性
plot_importance(model)
plt.show()